/*
 * ProbDensStat.h
 *
 *  Created on: Aug 25, 2010
 *      Author: bernardo
 */

#ifndef PROBDENSSTAT_H_
#define PROBDENSSTAT_H_

#include "ProbDens.hpp"

template<size_t nComp>
class ProbDensStat: public ProbDens<nComp> {
public:
	explicit ProbDensStat(size_t nX, double x0, double x1);

	explicit ProbDensStat(double dX);

	explicit ProbDensStat() {}

	void calcStats();

	double prob(int comp = -1) const;

	double meanX(int comp = -1) const;

	double varX(int comp = -1) const;

	void saveStat(std::ostream& out, double time=sqrt(-1)) const;

private:
	double prob_[nComp], meanX_[nComp], varX_[nComp];
};

template<size_t nComp>
ProbDensStat<nComp>::ProbDensStat(size_t nX, double x0, double x1) :
	ProbDens<nComp>(nX, x0, x1) {
}

template<size_t nComp>
ProbDensStat<nComp>::ProbDensStat(double dX) :
	ProbDens<nComp>(dX) {
}

template<size_t nComp>
double ProbDensStat<nComp>::prob(int comp) const {
	if (comp>=0) return prob_[comp];
	double s=0;
	for (size_t i=0;i!=nComp; i++) s += prob_[i];
	return s;
}

template<size_t nComp>
double ProbDensStat<nComp>::meanX(int comp) const {
	if (comp>=0) return meanX_[comp];
	double s=0;
	for (size_t i=0;i!=nComp; i++) s += meanX_[i]*prob_[i];
	return s/prob(-1);
}

template<size_t nComp>
double ProbDensStat<nComp>::varX(int comp) const {
	if (comp>=0) return varX_[comp];
	double s=0;
	for (size_t i=0;i!=nComp; i++) s += varX_[i]*prob_[i];
	return s/prob(-1);
}

template<size_t nComp>
void ProbDensStat<nComp>::calcStats() {
	for (size_t c = 0; c != nComp; c++) {
		double sP = 0, siP = 0, siiP = 0;
		double i = 0; //Avoid two conversions from int to double
		for (typename ProbDens<nComp>::iterator p=ProbDens<nComp>::begin(c),
				pend = ProbDens<nComp>::end(c); p != pend; p++, i++) {
			double pi = *p;
			sP += pi;
			siP += (pi *= i);
			siiP += (pi *= i);
		}
		double dx=ProbDens<nComp>::getDX(), x0=ProbDens<nComp>::getX(0);
		prob_[c] = sP * dx;
		meanX_[c] = x0 + dx * (siP+0.5*sP) / sP;
		varX_[c] = dx * dx * (siiP - siP * siP / sP) / sP;
	}
}

template<size_t nComp>
void ProbDensStat<nComp>::saveStat(std::ostream& out, double time) const {
	//if time=nan print the header
	if (time != time) {
		out	<< "time";
        for (int i = nComp==1 ? 0 : -1; i != nComp; i++)
			out << "\tint" << i << "\tmeanX" << i << "\tvarX" << i;
	}
	else {
		out << time ;
		for (int i = nComp==1 ? 0 : -1; i != nComp; i++)
			out << "\t" << prob(i) << "\t" << meanX(i) << "\t" << varX(i);
	}
	out << "\n";
}

#endif /* PROBDENSSTAT_H_ */
